# using DelimitedFiles
using PyCall
@pyimport numpy as np

function scaling(X)
    u,l=max(X...),min(X...)
    convert.(UInt8,round.((X .-l) ./(u-l) .*255))
end
function gen_data(shift_x,shift_y)
    x=range(-5,length=64,stop=5)
    zp=@. -((x + shift_x)^2 +(x +shift_y)'^2)/ 7.5
    zm=@. -((x - shift_x)^2 +(x -shift_y)'^2)/ 7.5
    z1=@. exp(zp) * exp(1im * zp)
    z2=@. exp(zm) * exp(1im * zm)
    zc=z1 + z2
    Ic=zc .* conj.(zc)
    Zt=@. zc + sqrt(z1 *conj(z1) + z2 *conj(z2))
    Abs=real.(Zt .* conj.(Zt)) |>scaling
    rZt=real.(Zt) |>scaling
    iZt=abs.(imag.(Zt)) |>scaling
    Abs,rZt,iZt
end

indata=zeros(UInt8,32*32,1,64,64)
outdata=zeros(UInt8,32*32,2,64,64)
cnt=1
for shift_x=range(-3,length=32,stop=3),shift_y=range(-3,length=32,stop=3)
    in,out1,out2=gen_data(shift_x,shift_y)
    indata[cnt,1,:,:] .=in
    outdata[cnt,1,:,:] .=out1
    outdata[cnt,2,:,:] .=out2
    global cnt+=1
end

np.save("indata.npy",indata)
np.save("outdata.npy",outdata)

# indata=zeros(UInt8,16*16,1,64,64)
# outdata=zeros(UInt8,16*16,2,64,64)
# shift_xy=randn(16*16,2)
# for id=1:256
#     in,out1,out2=gen_data(shift_xy[id,:]...)
#     indata[id,1,:,:] .=in
#     outdata[id,1,:,:] .=out1
#     outdata[id,2,:,:] .=out2
# end
#
# np.save("indata_t.npy",indata)
# np.save("outdata_t.npy",outdata)
